#include "Hooker.h"

#define EASY_SUCCESS(s) (NT_SUCCESS(status = s))
#define EASY_FAILED(s)	(!EASY_SUCCESS(s))

Hooker::Hooker() {
	ACLEntries_[0] = 0;
}

Hooker::~Hooker() {
	destroy();
}

void Hooker::destroy() {
	if(hookHandlesMap_.size() > 0) {
		//LhUninstallAllHooks();
		for(auto iter = hookHandlesMap_.begin(); iter != hookHandlesMap_.end(); ++iter) {
			LhUninstallHook(iter->second);
			delete iter->second;
			//iter->second = NULL;
			//iter = hookHandlesMap_.erase(iter);
		}
		hookHandlesMap_.clear();

		LhWaitForPendingRemovals();
	}

	for(auto const& mh : modulesHandles_) {
		FreeLibrary(mh.second);
	}
	modulesHandles_.clear();
}

std::pair<bool,NTSTATUS> Hooker::installHook(std::string const& moduleName,
									 std::string const& procName,
									 void* hookProc,
									 bool loadModule) {
	HMODULE hModule = (loadModule) ? (loadModule_(moduleName)) :
									 (GetModuleHandle(moduleName.c_str()));
	if(hModule == NULL) {
		return std::make_pair(false,0);
	}
	TRACED_HOOK_HANDLE hHook = new HOOK_TRACE_INFO();
	NTSTATUS status;
	if(EASY_FAILED(LhInstallHook(
			GetProcAddress(hModule,procName.c_str()),
			hookProc,
			(PVOID)0x12345678,
			hHook))) {
		delete hHook;
		return std::make_pair(false,status);
	}
	auto iter = hookHandlesMap_.find(procName);
	if(iter != hookHandlesMap_.end()) {
		//assert(false);
		return std::make_pair(false,0); // ,STATUS_SUCCESS);
	}
	hookHandlesMap_[procName] = hHook;
	return std::make_pair(true,0);
}

std::pair<bool,NTSTATUS> Hooker::activateHook(std::string const& procName,bool exclusive) {
	auto iter = hookHandlesMap_.find(procName);
	if(iter == hookHandlesMap_.end()) {
		return std::make_pair(false,0);
	}
	NTSTATUS status;
	if(EASY_FAILED((exclusive) ?
			(LhSetExclusiveACL(ACLEntries_,1,iter->second)) :
			(LhSetInclusiveACL(ACLEntries_,1,iter->second)))) {
		return std::make_pair(false,status);
	}
	return std::make_pair(true,0);
}

////

HMODULE Hooker::loadModule_(std::string const& name) {
	// Call LoadLibrary once per module
	auto iter = modulesHandles_.find(name);
	HMODULE hModule;
	if(iter == modulesHandles_.end()) {
		hModule = LoadLibrary(name.c_str());
		modulesHandles_[name] = hModule;
	}
	else {
		hModule = iter->second;
	}
	return hModule;
}